"""
输入：标注数据
输出：每个动作以及它的参数和关系
"""

class ActionType:
    def __init__(self):
        self.action_type = "NoneType"
        self.atomic_action_interactivity = "NoneInteractivity"
        self.atomic_action_dynamism = "NoneDynamism"
        self.action_temporal = "NoneTemporal"
        self.atomic_action_target = "NoneTarget"
        self.atomic_action_scene_switchability = "NoneSceneSwitchability"
        self.atomic_action_presentability = "NonePresentability"
        self.atomic_action_transformability = "NoneTransformability"


    # def __init__(self, action_type, atomic_action_interactivity, atomic_action_dynamism, action_temporal, atomic_action_target, atomic_action_scene_switchability, atomic_action_presentability, atomic_action_transformability):
    #     self.action_type = action_type
    #     self.atomic_action_interactivity = atomic_action_interactivity
    #     self.atomic_action_dynamism = atomic_action_dynamism
    #     self.action_temporal = action_temporal
    #     self.atomic_action_target = atomic_action_target
    #     self.atomic_action_scene_switchability = atomic_action_scene_switchability
    #     self.atomic_action_presentability = atomic_action_presentability
    #     self.atomic_action_transformability = atomic_action_transformability

    # 拼装 event_type
    def get_event_type(self):
        if self.action_type == "abstract action" or self.action_type == "combination action":
            return self.action_type
        return self.action_type + ", " + self.atomic_action_interactivity + ", " + self.atomic_action_dynamism + ", " + self.action_temporal + ", " + self.atomic_action_target + ", " + self.atomic_action_scene_switchability + ", " + self.atomic_action_presentability + ", " + self.atomic_action_transformability

class ActionRelevance:
    def __init__(self):
        self.initiator_list = []
        self.recipient_list = []
        self.indirector_list = []
        self.locale = []
        self.direction = ""
        self.background = ""
        self.region_list = []
        self.target = ""
        self.env_item_list = []
        self.time_list = []
        self.weather_list = []
        self.season_list = []

    # def __init__(self, initiator_list: list, recipient_list: list, indirector_list: list, locale: list, direction: str, background: str, region_list: list, target: list, env_item_list: list, time_list: list, weather_list: list, season_list: list):
    #     self.initiator_list = initiator_list
    #     self.recipient_list = recipient_list
    #     self.indirector_list = indirector_list
    #     self.locale = locale
    #     self.direction = direction
    #     self.background = background
    #     self.region_list = region_list
    #     self.target = target
    #     self.env_item_list = env_item_list
    #     self.time_list = time_list
    #     self.weather_list = weather_list
    #     self.season_list = season_list

    # 用模板拼装 arg
    def get_arg_with_template(self, role, text, start, end):
        return {
            "role": role,
            "text": text,
            "span": [start, end]
        }

    # 拼装 args
    def get_args(self):
        arg_list = []
        if self.initiator_list:
            for initiator in self.initiator_list:
                arg_list.append(initiator)
        if self.recipient_list:
            for recipient in self.recipient_list:
                arg_list.append(recipient)
        if self.indirector_list:
            for indirector in self.indirector_list:
                arg_list.append(indirector)
        if self.locale:
            for loc in self.locale:
                arg_list.append(loc)
        if self.direction:
            arg_list.append(self.direction)
        if self.background:
            arg_list.append(self.background)
        if self.region_list:
            for region in self.region_list:
                arg_list.append(region)
        if self.target:
            arg_list.append(self.target)
        if self.env_item_list:
            for env_item in self.env_item_list:
                arg_list.append(env_item)
        if self.time_list:
            for time in self.time_list:
                arg_list.append(time)
        if self.weather_list:
            for weather in self.weather_list:
                arg_list.append(weather)
        if self.season_list:
            for season in self.season_list:
                arg_list.append(season)
        return arg_list

class MyAction:
    def __init__(self, action_type: ActionType, action_trigger_words, trigger_start, trigger_end, action_relevance: ActionRelevance):
        self.event_type = action_type.get_event_type()
        self.trigger = {
            "text": action_trigger_words,
            "span": [trigger_start, trigger_end]
        }
        self.args = action_relevance.get_args()

class ActionSeqRel:
    def __init__(self, relation, h_text, h_start, h_end, t_text, t_start, t_end):
        self.relation = relation
        self.head = {
            "text": h_text,
            "span": [
                h_start,
                h_end
            ]
        }
        self.tail = {
            "text": t_text,
            "span": [
                t_start,
                t_end
            ]
        }

    def to_dict(self):
        return {
            "relation": self.relation,
            "head": self.head,
            "tail": self.tail
        }

import os
import json
from pyvis.network import Network
import networkx as nx
from itertools import combinations
import matplotlib.pyplot as plt

cur_path = os.path.abspath(os.path.dirname(__file__))

# 从 json 文件中读数据
def read_from_json(filename):
    try:
        with open(os.path.join(cur_path, filename), 'r', encoding='utf-8') as f:
            data = json.load(f)
        return data
    except Exception as e:
        print("读取文件出错：", e)

# 写数据到 json 文件
def write_to_json(filename, data):

    # 获取文件所在的目录路径
    dir_path = os.path.dirname(os.path.join(cur_path, filename))

    # 如果文件夹不存在则先创建文件夹
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

    try:
        with open(os.path.join(cur_path, filename), 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=4)
    except Exception as e:
        print("写入文件出错：", e)

# 从一般文件中读数据
def read_from_file(filename):
    try:
        with open(os.path.join(cur_path, filename), 'r', encoding='utf-8') as f:
            data = f.readlines()
        return data
    except Exception as e:
        print("读取文件出错：", e)

# 向文件中写入数据
def add_to_file(filename, data):
    try:
        with open(os.path.join(cur_path, filename), 'a+', encoding='utf-8') as f:
            f.write(data)
    except Exception as e:
        print("写入文件出错：", e)

# 清空文件
def clear_file(filename):
    try:
        with open(os.path.join(cur_path, filename), 'w', encoding='utf-8') as f:
            f.write('')
    except Exception as e:
        print("清空文件出错：", e)

def validate_a_piece_of_story(story_name, label_data):

    print(f"\n---------- ### 正在验证 \"{story_name}\" 的数据 ### ----------")

    # 读入配置数据
    story = read_from_file(f'../story/{story_name}.txt')
    print(f"读取了 {len(story)} 行故事，共 {len(''.join(story))} 个字符")
    # list_action_label = read_from_json('../config/action_label.json')
    # print(f"读取了 {len(list_action_label)} 个动作标签")
    list_from_name = read_from_json('../config/from_name.json')
    print(f"读取了 {len(list_from_name)} 个动作子标签")
    list_rel = read_from_json('../config/action_rel.json')
    print(f"读取了 {len(list_rel)} 个动作关系标签")

    # 读入标注数据 1、2
    # 1 读事件标注数据
    annotate_result = None
    for data_ele in label_data:
        # 如果 data_ele["data"]["text"] 包含字符串 story_name 如 "小红帽"，并且 data_ele["file_upload"] 不以 "_for_rels.json" 结尾

        if story_name in data_ele["data"]["text"] and not data_ele["file_upload"].endswith("_for_rels.json"):
            annotate_result = data_ele["annotations"][0]["result"]
            break

    # 2 读时序关系标注数据
    annotate_result_for_rels = None
    for data_ele in label_data:
        # 如果 data_ele["data"]["text"] 包含字符串 story_name 如 "小红帽"，并且 data_ele["file_upload"] 以 "_for_rels.json" 结尾
        if story_name in data_ele["data"]["text"] and data_ele["file_upload"].endswith("_for_rels.json"):
            annotate_result_for_rels = data_ele["annotations"][0]["result"]
            break
    
    # 关键动作映射表：键是 id，值是 MyAction 对象
    action_dict = {}

    # 实体映射表：键是 id，值是 start, end 组成的字典
    ent_dict = {}

    # 时序关系列表
    action_seq_rel_list = []

    # 遍历 annotate_result 列表，存储实体对应于原文的位置
    for result_ele in annotate_result:
        if "type" in result_ele and result_ele["type"] == "labels":
            if result_ele["id"] not in ent_dict:
                ent_dict[result_ele["id"]] = {
                    "label": result_ele["value"]["labels"][0],
                    "start": result_ele["value"]["start"],
                    "end": result_ele["value"]["end"]
                }

    # 遍历 result_little_red_hat 列表，封装动作
    # 1. 封装 ActionType；2. 封装 ActionRelevance 3. 拼接 MyAction
    # 1
    if None != annotate_result:
        action_type_dict = {}
        for result_ele in annotate_result:
            # 拼装 event_type：如果 result_ele 有 "value" 键，并且 "value" 键有 "labels" 键，并且 from_name 在 list_from_name 中
            if "value" in result_ele and "labels" in result_ele["value"] and "from_name" in result_ele and result_ele["from_name"] in list_from_name and result_ele["from_name"] != "Label_Action_Relevance":

                # span 的 start 大于等于 story 的长度的，暂时跳过不处理
                if result_ele["value"]["start"] >= len(story[0]):
                    continue

                if result_ele["id"] not in action_type_dict:
                    action_type_dict[result_ele["id"]] = ActionType()

                # 按照 from_name 赋值 MyAction 属性
                if result_ele["from_name"] == "Label_Action_Type":
                    action_type_dict[result_ele["id"]].action_type = result_ele["value"]["labels"][0]
                elif result_ele["from_name"] == "Label_Atomic_Action_Interactivity":
                    action_type_dict[result_ele["id"]].atomic_action_interactivity = result_ele["value"]["labels"][0]
                elif result_ele["from_name"] == "Label_Atomic_Action_Dynamism":
                    action_type_dict[result_ele["id"]].atomic_action_dynamism = result_ele["value"]["labels"][0]
                elif result_ele["from_name"] == "Label_Action_Temporal":
                    action_type_dict[result_ele["id"]].action_temporal = result_ele["value"]["labels"][0]
                elif result_ele["from_name"] == "Label_Atomic_Action_Target":
                    action_type_dict[result_ele["id"]].atomic_action_target = result_ele["value"]["labels"][0]
                elif result_ele["from_name"] == "Label_Atomic_Action_Scene_Switchability":
                    action_type_dict[result_ele["id"]].atomic_action_scene_switchability = result_ele["value"]["labels"][0]
                elif result_ele["from_name"] == "Label_Atomic_Action_Presentability":
                    action_type_dict[result_ele["id"]].atomic_action_presentability = result_ele["value"]["labels"][0]
                elif result_ele["from_name"] == "Label_Atomic_Action_Transformability":
                    action_type_dict[result_ele["id"]].atomic_action_transformability = result_ele["value"]["labels"][0]
        
        # 2
        action_relevance_dict = {}
        for result_ele in annotate_result:
            # 如果 result_ele["type"] 是 "relation"，result_ele["to_id"] 在 action_type_dict 的键中，则遍历 action_type_dict，向 action_relevance_dict 添加对应的 id 和 action_relevance
            if "type" in result_ele and result_ele["type"] == "relation" and "to_id" in result_ele and result_ele["to_id"] in action_type_dict.keys():
                # 暂时忽略 Modifier 非实体
                if "Modifier" == ent_dict[result_ele["from_id"]]["label"]:
                    continue
                if result_ele["to_id"] not in action_relevance_dict:
                    action_relevance_dict[result_ele["to_id"]] = ActionRelevance()
                t_start = ent_dict[result_ele["from_id"]]["start"]
                t_end = ent_dict[result_ele["from_id"]]["end"]
                # 遍历关系标签列表
                for a_rel in result_ele["labels"]:
                    if a_rel in list_rel:
                        if "act:initiator_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].initiator_list.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Initiator", story[0][t_start:t_end], t_start, t_end))
                        elif "act:recipient_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].recipient_list.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Recipient", story[0][t_start:t_end], t_start, t_end))
                        elif "act:indirector_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].indirector_list.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Indirector", story[0][t_start:t_end], t_start, t_end))
                        elif "act:locale_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].locale.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Locale", story[0][t_start:t_end], t_start, t_end))
                        elif "act:direction_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].direction = action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Direction", story[0][t_start:t_end], t_start, t_end)
                        elif "act:background_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].background = action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Background", story[0][t_start:t_end], t_start, t_end)
                        elif "act:region_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].region_list.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Region", story[0][t_start:t_end], t_start, t_end))
                        elif "act:target_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].target = action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Target", story[0][t_start:t_end], t_start, t_end)
                        elif "act:env_item_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].env_item_list.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Env_item", story[0][t_start:t_end], t_start, t_end))
                        elif "act:time_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].time_list.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Time", story[0][t_start:t_end], t_start, t_end))
                        elif "act:weather_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].weather_list.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Weather", story[0][t_start:t_end], t_start, t_end))
                        elif "act:season_of" == a_rel:
                            action_relevance_dict[result_ele["to_id"]].season_list.append(action_relevance_dict[result_ele["to_id"]].get_arg_with_template("Season", story[0][t_start:t_end], t_start, t_end))

        # 3
        # 遍历 action_type_dict
        # debug
        if "白雪公主" == story_name:
            if "J7mATR8sr2" not in ent_dict.keys():
                print("白雪公主的 J7mATR8sr2 动作没有在实体列表中找到")
        for action_id in action_type_dict.keys():
            trigger_start = ent_dict[action_id]["start"]
            trigger_end = ent_dict[action_id]["end"]
            trigger_words = story[0][trigger_start:trigger_end]
            action_dict[action_id] = MyAction(action_type=action_type_dict[action_id], action_trigger_words=trigger_words, trigger_start=trigger_start, trigger_end=trigger_end, action_relevance=action_relevance_dict[action_id])

        # 初始化动作文件
        # clear_file(f"../out/{story_name}/actions.json")

        # 初始化动作关系文件
        # clear_file(f"../out/{story_name}/action_seq_rel.json")

        # 将 action_dict 按照 value["trigger"]["span"][0] 排序
        # key_action_list = sorted(((k, v.__dict__) for k, v in action_dict.items()), key=lambda item: item[1]['trigger']['span'][0])
        key_action_list = sorted(((v.__dict__) for v in action_dict.values()), key=lambda item: item['trigger']['span'][0])

        # 动作列表写入文件
        write_to_json(f'../out/{story_name}/actions.json', key_action_list)

    # 遍历 annotate_result_for_rels 列表，封装时序关系列表
    if None != annotate_result_for_rels:
        G = nx.DiGraph() # 创建一个 networkx 图
        added_nodes = set() # 用来检查节点在图中是否被添加过
        my_nodes = [] # 用来存储节点信息
        my_edges = [] # 用来存储边的信息
        for result_ele in annotate_result_for_rels:
            if "type" in result_ele and result_ele["type"] == "relation":
                # 如果 result_ele["labels"][0] 以 "seq:" 开头
                if result_ele["labels"][0].startswith("seq:"):
                    # relation 的值是 "seq:" 之后的部分
                    relation = result_ele["labels"][0][4:]
                    h_start = ent_dict[result_ele["from_id"]]["start"]
                    h_end = ent_dict[result_ele["from_id"]]["end"]
                    h_text = story[0][h_start:h_end]
                    t_start = ent_dict[result_ele["to_id"]]["start"]
                    t_end = ent_dict[result_ele["to_id"]]["end"]
                    t_text = story[0][t_start:t_end]
                    action_seq_rel_list.append(ActionSeqRel(relation=relation, h_text=h_text, h_start=h_start, h_end=h_end, t_text=t_text, t_start=t_start, t_end=t_end).to_dict())

                    
                    ### ---------- 将节点添加到 nx.Graph 的实例 G 中，借助 networkx 的特性来确定节点的位置 ----------
                    # 节点标记
                    from_label = f"{[d['text'] for d in action_dict[result_ele['from_id']].args if d['role'] == 'Initiator']}, {story[0][h_start:h_end]}:[{h_start},{h_end})"
                    to_label = f"{[d['text'] for d in action_dict[result_ele['to_id']].args if d['role'] == 'Initiator']}, {story[0][t_start:t_end]}:[{t_start},{t_end})"
                    # 添加节点
                    if result_ele["from_id"] not in added_nodes:
                        G.add_node(from_label)
                        G.nodes[from_label]['group'] = f"{[d['text'] for d in action_dict[result_ele['from_id']].args if d['role'] == 'Initiator']}".lower()
                        my_nodes.append(from_label)
                    if result_ele["to_id"] not in added_nodes:
                        G.add_node(to_label)
                        G.nodes[to_label]['group'] = f"{[d['text'] for d in action_dict[result_ele['to_id']].args if d['role'] == 'Initiator']}".lower()
                        my_nodes.append(to_label)

                    added_nodes.add(result_ele["from_id"])
                    added_nodes.add(result_ele["to_id"])
                    
                    # 添加边
                    # print(f"debug:添加边 {from_label} -> {to_label}，类型为 {relation}")
                    G.add_edge(from_label, to_label, type=relation)
                    my_edges.append((from_label, to_label, relation))
                    ### ---------- 将节点添加到 nx.Graph 的实例 G 中，借助 networkx 的特性来确定节点的位置 ----------
        print("-------------------------------------------------------------")

        # 动作关系列表写入文件
        write_to_json(f'../out/{story_name}/action_seq_rel.json', action_seq_rel_list)

        # 用 pyvis 可视化时序关系图
        # 1 创建一个 pyvis 网络图
        net = Network(heading=f"{story_name} 的动作时序关系图")
        # 2 设置网络图的布局
        net.set_options("""
        var options = {
            "nodes": {
                "size": 100,
                "font": {
                    "size": 58,
                    "face": "LXGW Wenkai Mono"
                }
            },
            "edges": {
                "color": {
                    "inherit": true,
                    "opacity": 1
                },
                "smooth": {
                    "type": "straight",
                    "forceDirection": "none"
                }
            },
            "physics": {
                "enabled": false
            }
        }
        """)
        # 3 生成网络图

        # 获取节点到组别的映射
        groups = {}
        for a_node, attr in G.nodes(data=True):
            group = attr.get('group', None)
            if group is not None:
                if group not in groups:
                    groups[group] = []
                groups[group].append(a_node)
        
        # 初始化布局字典
        pos = {}
        
        # 为每个组分配列位置
        for i, (group, nodes) in enumerate(groups.items()):
            for j, node in enumerate(nodes):
                # 按照列排列节点，这里使用组的索引作为 x 坐标，节点在组内的索引作为 y 坐标，分别缩放让显示效果更好
                pos[node] = (i * 1500, j * 300)

        # 添加节点到网络图

        for node, attr in G.nodes(data=True):
            net.add_node(node, label=node, x=pos[node][0], y=pos[node][1], group=attr.get('group', None))
        for node1, node2, edge_type in G.edges(data=True):
            if "before" == edge_type["type"]:
                net.add_edge(node1, node2, arrows={'to': {'enabled': True, 'scaleFactor': 2}}, color="#3c52fd")
            elif "sync" == edge_type["type"]:
                net.add_edge(node1, node2, color="#13e807")

        # 4 保存到 html 文件，显示网络图
        if "保存到 html 文件":
            output_html_path = os.path.join(cur_path, f"../out/{story_name}/action_seq_rel.html")
            net.show(output_html_path, notebook=False)
            # 修复 html 中显示两次标题的 bug
            import re
            # 读取HTML文件内容
            with open(output_html_path, 'r', encoding='utf-8') as file:
                html_content = file.read()

            # 定义要删除的内容的正则表达式
            pattern = r'<center>\s*<h1>.*?</h1>\s*</center>'

            # 使用正则表达式替换掉第一个匹配的内容
            new_html_content = re.sub(pattern, '', html_content, count=1)

            # 将修改后的内容写回文件
            with open(output_html_path, 'w', encoding='utf-8') as file:
                file.write(new_html_content)

    # print(f"---------- ### \"{story_name}\" 的数据验证完成。### ----------\n")

if __name__ == '__main__':
    
    # 读数据
    json_label_data = read_from_json('../project-6-at-2024-12-25-20-35-7acb15df-multi-stories.json')
    print(f"读取了 {len(json_label_data)} 行标注数据")

    # 验证
    validate_a_piece_of_story("小红帽", json_label_data)
    validate_a_piece_of_story("白雪公主", json_label_data)
    validate_a_piece_of_story("青蛙王子", json_label_data)
    validate_a_piece_of_story("龟兔赛跑", json_label_data)